import streamlit as st
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

def plot_price_drop_bar(data, title="股票跌幅排名"):
    """
    绘制股票跌幅排名条形图
    
    Args:
        data: 包含股票数据的DataFrame
        title: 图表标题
    
    Returns:
        plotly图表对象
    """
    # 确保数据按跌幅排序
    sorted_data = data.sort_values('drop_percent', ascending=True)
    
    # 添加股票名称到标签
    sorted_data['label'] = sorted_data['ts_code'] + ' - ' + sorted_data['name']
    
    fig = px.bar(
        sorted_data,
        y='label',
        x='drop_percent',
        orientation='h',
        title=title,
        labels={'drop_percent': '跌幅(%)', 'label': '股票'},
        color='drop_percent',
        color_continuous_scale='Reds',
        text=sorted_data['drop_percent'].round(2).astype(str) + '%'
    )
    
    fig.update_traces(textposition='outside')
    fig.update_layout(height=500)
    
    return fig

def plot_industry_distribution(data, title="行业分布"):
    """
    绘制行业分布饼图
    
    Args:
        data: 包含股票数据的DataFrame
        title: 图表标题
    
    Returns:
        plotly图表对象
    """
    industry_counts = data['industry'].value_counts().reset_index()
    industry_counts.columns = ['industry', 'count']
    
    fig = px.pie(
        industry_counts,
        values='count',
        names='industry',
        title=title,
        hole=0.4
    )
    
    fig.update_layout(height=500)
    
    return fig

def plot_volume_drop_scatter(data, title="成交量与跌幅关系"):
    """
    绘制成交量与跌幅关系散点图
    
    Args:
        data: 包含股票数据的DataFrame
        title: 图表标题
    
    Returns:
        plotly图表对象
    """
    fig = px.scatter(
        data,
        x='vol',
        y='drop_percent',
        size='amount',
        color='industry',
        hover_name='name',
        hover_data=['ts_code', 'open', 'close'],
        title=title,
        labels={
            'vol': '成交量', 
            'drop_percent': '跌幅(%)', 
            'amount': '成交额', 
            'industry': '行业',
            'name': '股票名称',
            'ts_code': '股票代码',
            'open': '开盘价',
            'close': '收盘价'
        }
    )
    
    fig.update_layout(height=600)
    
    return fig

def plot_market_overview(market_data):
    """
    绘制市场概览图表
    
    Args:
        market_data: 市场概览数据字典
    
    Returns:
        plotly图表对象
    """
    # 创建子图
    fig = make_subplots(
        rows=1, cols=2,
        specs=[[{"type": "indicator"}, {"type": "pie"}]],
        subplot_titles=("上证指数变化", "股票涨跌分布")
    )
    
    # 添加指数变化指示器
    fig.add_trace(
        go.Indicator(
            mode="delta",
            value=market_data['index_change_percent'],
            delta={
                "reference": 0,
                "valueformat": ".2f",
                "suffix": "%",
                "increasing": {"color": "red"},
                "decreasing": {"color": "green"}
            },
            title={
                "text": f"上证指数 ({market_data['trade_date']})"
            }
        ),
        row=1, col=1
    )
    
    # 添加涨跌分布饼图
    labels = ['上涨', '下跌', '平盘']
    values = [market_data['up_stocks'], market_data['down_stocks'], market_data['flat_stocks']]
    colors = ['red', 'green', 'gray']
    
    fig.add_trace(
        go.Pie(
            labels=labels,
            values=values,
            marker=dict(colors=colors),
            textinfo='label+percent',
            hole=0.3
        ),
        row=1, col=2
    )
    
    fig.update_layout(
        height=400,
        showlegend=False
    )
    
    return fig

def plot_stock_detail(stock_data, historical_data=None):
    """
    绘制单只股票的详细信息图表
    
    Args:
        stock_data: 单只股票的数据Series
        historical_data: 股票的历史数据DataFrame (可选)
    
    Returns:
        plotly图表对象
    """
    # 创建图表
    fig = go.Figure()
    
    # 添加OHLC信息
    fig.add_trace(
        go.Indicator(
            mode="number+delta",
            value=stock_data['close'],
            delta={
                "reference": stock_data['open'],
                "valueformat": ".2f",
                "increasing": {"color": "red"},
                "decreasing": {"color": "green"}
            },
            title={
                "text": f"{stock_data['name']} ({stock_data['ts_code']})"
            },
            domain={'row': 0, 'column': 0}
        )
    )
    
    # 如果有历史数据，添加K线图
    if historical_data is not None and not historical_data.empty:
        fig.add_trace(
            go.Candlestick(
                x=historical_data.index,
                open=historical_data['open'],
                high=historical_data['high'],
                low=historical_data['low'],
                close=historical_data['close'],
                increasing_line_color='red',
                decreasing_line_color='green'
            )
        )
        
        fig.update_layout(
            xaxis_title='日期',
            yaxis_title='价格'
        )
    
    fig.update_layout(height=500)
    
    return fig
